from copy import deepcopy
from math import log
from random import randrange
import time
from bayesNet import PT, BN
from mle import mleFromCompleteData
import networks_repository

__author__ = 'Matija'

"""
Greedy hill climbing with random restarts algorithm. This is an extension of a basic hill climbing with random restarts
that works on bayes networks.

The basic algorithm idea is to create N-times a random network, run hill climbing algorithm on this network and return
a network with the maximum score.

The hill climbing idea is to every time find all possible neighbours of the given network. This means that for every
node pair are 3 possibilities: no edge, edge from node one to node two, edge from node two to node one.
This means that are 3 * N (N is a number of pairs) possible combinations.

The score calculation idea is:
- score(M|D) = LL(M|D) - psi(N) * |G|
- |G| = sum( (#Xi - 1) * #Ui )
    - Xi = variables
    - Ui = all parents for variables
- LL(M|D) = log L(M|D) = LL (Vml|D)
- L(V|D) = pi( Pm,v(di) )


"""

# returns all nodes that have an edge from node
def allNodesWithEdgeFromNode(node, edges):
    e = list()
    for edge in edges:
        if edge[0] == node:
            e.append(edge[1])
    return e

# checks if node has any incoming edges
def hasIncomingNodes(parents, node, edges):
    if not parents:
        return False

    for parent in parents:
        if (parent, node) in edges:
            return True

    return False

def areAllNodesConnected(network):
    # check if all nodes has at least one incoming or outgoing connection
    nodes = network.nodes
    for node in nodes:
        parents = network.parents(node)

        # if no incoming connections, check outgoing
        if not parents:
            hasOut = False
            for node1 in nodes:
                if node1 != node:
                    p = network.parents(node1)
                    if node in p:
                        hasOut = True
                        break
            if not hasOut:
                return False
    return True

def isNetworkValid(network):
    nodes = network.nodes       # all graph nodes
    edges = list()              # list of edges
    q = list()                  # all nodes with no parents

    # load all nodes without parents into q
    for node in nodes:
        parents = network.parents(node)
        if not parents:
            q.append(node)

    # transform network into edges
    for node in nodes:
        parents = network.parents(node)
        for parent in parents:
            edges.append( (parent, node) )

    while q:
        n = q.pop()      # remove node from list
        for m in allNodesWithEdgeFromNode(n, edges):
            edges.remove( (n, m) )      # remove edge from n to m
            parents = network.parents(m)

            # if m has no incoming edges put into q
            if not hasIncomingNodes(parents, m, edges):
                q.append(m)

    if edges:
        return False
    else:
        return True

def calcLL(network, exampleTable):
    LL = 0
    # go through all examples in exampleTable and for each example calculate probability
    for example in exampleTable.data:
        probability = 1
        # go through all nodes in network and find example match in PT
        # PT probabilities are multiplied
        for node in network.nodes:
            # for each node get a PT and find vales for nodes in PT
            pt = network.probabilities[node]
            vals = list()
            for pnode in pt.nodes:
                vals.append(example.values[exampleTable.vars_to_indices[pnode]])

            # find a matching line in PT
            for value in pt.probabilityTable:
                if value[0:len(vals)] == vals:
                    probability *= value[len(value) - 1]
                    break
        # logarithm
        LL += log(probability, 2)
    return LL

def calcPsi(n, nodes):
    # x * log(n) with base 2
#    return log(n, 2) * pow(nodes, 2) / 200
    return log(n, 2) / 2

def calcG(network):
    # |G| = sum ( |Xi * Ui| ); |Xi * Ui| = (Xi# - 1) * Ui#
    # Xi = variables, Ui = variables parents

    # comment: X = [True, False] --> Xi# = 2 --> (Xi# - 1 ) = 1; Ui# = 2 ^ noOfParents
    G = 0
    nodes = network.nodes
    for node in nodes:
        parents = network.parents(node)
        G += pow(2, len(parents))
    return G

def scoreNetwork(network, exampleTable):
    # calc network probability tables
    mleFromCompleteData(network, exampleTable)

    # calc score
    LL = calcLL(network, exampleTable)
    psi = calcPsi(len(exampleTable), len(network.nodes))
    G = calcG(network)
    score = LL - psi * G
    return score

def findNetworkNeighbours(network):
    neighbours = []
    nodes = network.nodes

    # get all node pairs - for each pair will be generated 2 additional networks
    # there are 3 possibilities: no connection, connection left, connection right
    allPairs = []
    for node1 in nodes:
        for node2 in nodes:
            if node1 != node2 and (node1, node2) not in allPairs and (node2, node1) not in allPairs:
                allPairs.append( (node1, node2) )

    # iterate over all node pairs and for each pair change one edge
    for pair in allPairs:
        neighbourPair = generateNeighbourPair(network, pair)
        if isNetworkValid(neighbourPair[0]):
            neighbours.append(neighbourPair[0])
        if isNetworkValid(neighbourPair[1]):
            neighbours.append(neighbourPair[1])

    return neighbours

def generateNeighbourPair(network, pair):
    edges = list()
    nodes = network.nodes

    # transform network into edges
    for node in nodes:
        parents = network.parents(node)
        for parent in parents:
            edges.append( (parent, node) )

    # generate neighbours
    edges1 = deepcopy(edges)
    edges2 = deepcopy(edges)
    if pair in edges:
        # in case the network contains this edge
        edges1.remove(pair)
        edges2.remove(pair)
        edges2.append( (pair[1], pair[0]) )
    elif (pair[1], pair[0]) in edges:
        # in case the network contains reverted edge
        edges1.remove( (pair[1], pair[0]))
        edges2.remove( (pair[1], pair[0]))
        edges2.append(pair)
    elif pair not in edges:
        # in case the network doesn't contain this edge
        edges1.append(pair)
        edges2.append( (pair[1], pair[0]) )

    # transform into BN
    pts = dict()
    for node in nodes:
        parents = getParents(node, edges1)
        parents.append(node)
        pts[node] = PT(parents, [])
    neighbour1 = BN(nodes, pts)

    pts = dict()
    for node in nodes:
        parents = getParents(node, edges2)
        parents.append(node)
        pts[node] = PT(parents, [])

    neighbour2 = BN(nodes, pts)

    return neighbour1, neighbour2

def hillClimbingAlgorithm(network, exampleTable):
    # set best network to given (random) network and best score to this network score
    bestNetwork = network
    bestScore = scoreNetwork(bestNetwork, exampleTable)
    generated = 0

    # run while network can be improved
    while True:
        # find all neighbours for bestNetwork
        neighbours = findNetworkNeighbours(bestNetwork)
        currentBestNetwork = None
        currentBestScore = float("-inf")
        generated += len(neighbours)

        # for all neighbours compute theirs score and find the best network
        for neighbour in neighbours:
            neighbourScore = scoreNetwork(neighbour, exampleTable)
            if neighbourScore > currentBestScore:
                currentBestScore = neighbourScore
                currentBestNetwork = neighbour

        # if none of neighbours is best than current best network then the algorithm has found the local maxima
        # else set best neighbour to best network and search again
        if currentBestScore <= bestScore:
            break
        else:
            bestScore = currentBestScore
            bestNetwork = currentBestNetwork

    return bestNetwork, bestScore, generated

def createRandomNetwork(nodes):
    edges = list()
    max = 10000
    # create edges from all pairs
    for node1 in nodes:
        for node2 in nodes:
            if node1 != node2 and (node1, node2) not in edges and (node2, node1) not in edges:
                randomNumber = randrange(0, max, 1)
                if randomNumber > (max * 0.33) and randomNumber < (max * 0.66):
                    edges.append((node1, node2))
                elif randomNumber >= (max * 0.66):
                    edges.append((node2, node1))

    pts = dict()
    for node in nodes:
        parents = getParents(node, edges)
        parents.append(node)
        pts[node] = PT(parents, [])

    return BN(nodes, pts)

# get all parents for node from edges
def getParents(node, edges):
    parents = list()
    for edge in edges:
        if edge[1] == node:
            parents.append(edge[0])
    return parents

def networkToPng(network, outputName):
    network.to_png(outputName)

def greedyHillClimbingWithRandomRestarts(nodes, exampleTable, outputName, printScore = True, printNetwork = True):
    # params
    numberOfRandomRestarts = 50
    currentBestScore = float("-inf")
    currentBestNetwork = None
    generatedNetworks = 0

    # start hill climbing with random networks
    for restart in range(0, numberOfRandomRestarts):
        while True:
            randomNetwork = createRandomNetwork(nodes)
            if isNetworkValid(randomNetwork):
                break
        

        # hill climbing algorithm - returns best network and its score
        bestNetwork, bestScore, generated = hillClimbingAlgorithm(randomNetwork, exampleTable)
        generatedNetworks += generated
        if bestScore > currentBestScore:
            currentBestScore = bestScore
            currentBestNetwork = bestNetwork
    if printNetwork:
        networkToPng(currentBestNetwork, outputName)
    if printScore:
        print 'Best network saved as ' + outputName + '.png'
        print 'Best network score: ', currentBestScore
        print 'Algorithm generated %s networks' % (generatedNetworks)
    return currentBestScore, generatedNetworks


def runTests():
    start = time.clock()
    nodes, data = networks_repository.smallNetwork()
    greedyHillClimbingWithRandomRestarts(nodes, data, 'small')
    elapsed = time.clock() - start
    print 'Small :: Elapsed time in seconds:',
    print elapsed

#    start = time.clock()
#    nodes, data = networks_repository.asiaSmallNetwork()
#    greedyHillClimbingWithRandomRestarts(nodes, data, 'asiaSmall')
#    elapsed = time.clock() - start
#    print 'Asia small :: Elapsed time in seconds: ',
#    print elapsed

    start = time.clock()
    nodes, data = networks_repository.coronaryNetwork()
    greedyHillClimbingWithRandomRestarts(nodes, data, 'coronary')
    elapsed = time.clock() - start
    print 'Coronary :: Elapsed time in seconds: ',
    print elapsed

    start = time.clock()
    nodes, data = networks_repository.asiaBigNetwork()
    greedyHillClimbingWithRandomRestarts(nodes, data, 'asiaBig')
    elapsed = time.clock() - start
    print 'Asia big :: Elapsed time in seconds: ',
    print elapsed

    start = time.clock()
    nodes, data = networks_repository.winterNetwork()
    greedyHillClimbingWithRandomRestarts(nodes, data, 'winter')
    elapsed = time.clock() - start
    print 'Winter :: Elapsed time in seconds: ',
    print elapsed

def runCoronaryTests():
    print ''
    print ''
    print '------------------------------ Starting coronary tests ----------------------------------'
    exampleSizes = [200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800]
    for x in exampleSizes:
        print''
        print 'Starting: %s' % (time.time())
        start = time.time()
        nodes, data = networks_repository.coronaryNetwork(x)
        scoreAvg = 0
        generatedAvg = 0
        for i in range(0, 10):
            name = 'coronary_%s_iter_%s' % (x, i)
            score, generatedNetworks = greedyHillClimbingWithRandomRestarts(nodes, data, name, False, True)
            scoreAvg += score
            generatedAvg += generatedNetworks
        scoreAvg /= 10
        normalizedScore = scoreAvg / x
        generatedAvg /= 10
        print 'Avg score for coronary exampleSet size %s : %s ' % (x, scoreAvg)
        print 'Avg normalized score for coronary exampleSet size %s : %s ' % (x, normalizedScore)
        print 'Avg generated for coronary exampleSet size %s : %s ' % (x, generatedAvg)
        elapsed = (time.time() - start) / 10
        print 'Avg time consumption for algorithm with exampleSet size %s: Elapsed: %s in seconds' % (x, elapsed)
        print 'Stoping: %s' %(time.time())
        print ''

def runWinterTests():
    print ''
    print ''
    print '------------------------------ Starting winter tests ----------------------------------'
    exampleSizes = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]
    for x in exampleSizes:
        print''
        print 'Starting: %s' % (time.time())
        start = time.time()
        nodes, data = networks_repository.winterNetwork(x)
        scoreAvg = 0
        generatedAvg = 0
        for i in range(0, 10):
            name = 'winter_%s_iter_%s' % (x, i)
            score, generatedNetworks = greedyHillClimbingWithRandomRestarts(nodes, data, name, False, True)
            scoreAvg += score
            generatedAvg += generatedNetworks
        scoreAvg /= 10
        normalizedScore = scoreAvg / x
        generatedAvg /= 10
        print 'Avg score for winter exampleSet size %s : %s ' % (x, scoreAvg)
        print 'Avg normalized score for winter exampleSet size %s : %s ' % (x, normalizedScore)
        print 'Avg generated for winter exampleSet size %s : %s ' % (x, generatedAvg)
        elapsed = (time.time() - start) / 10
        print 'Avg time consumption for algorithm with exampleSet size %s: Elapsed: %s in seconds' % (x, elapsed)
        print 'Stoping: %s' %(time.time())
        print ''

def runAsiaTests():
    print ''
    print ''
    print '------------------------------ Starting asia tests ----------------------------------'
    exampleSizes = [500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500, 5000]
    for x in exampleSizes:
        print''
        print 'Starting: %s' % (time.time())
        start = time.time()
        nodes, data = networks_repository.asiaBigNetwork(x)
        scoreAvg = 0
        generatedAvg = 0
        for i in range(0, 10):
            name = 'asia_%s_iter_%s' % (x, i)
            score, generatedNetworks = greedyHillClimbingWithRandomRestarts(nodes, data, name, False, True)
            scoreAvg += score
            generatedAvg += generatedNetworks
        scoreAvg /= 10
        normalizedScore = scoreAvg / x
        generatedAvg /= 10
        print 'Avg score for asia exampleSet size %s : %s ' % (x, scoreAvg)
        print 'Avg normalized score for asia exampleSet size %s : %s ' % (x, normalizedScore)
        print 'Avg generated for asia exampleSet size %s : %s ' % (x, generatedAvg)
        elapsed = (time.time() - start) / 10
        print 'Avg time consumption for algorithm with exampleSet size %s: Elapsed: %s in seconds' % (x, elapsed)
        print 'Stoping: %s' %(time.time())
        print ''

if __name__=="__main__":
    runCoronaryTests()
    runAsiaTests()
    runWinterTests()
    runTests()